from paddleocr import PaddleOCR
import cv2
import os
import logging
from threading import Thread
import json


class SyncPredict(Thread):
    
    def __init__(self, imgs, data):
        super().__init__()
        self.imgs = imgs
        self.data = data
        
    def run(self):
        res = self.ocr_predict(self.imgs)
        for info in res:
            new_s = info.replace(' ', '')
            new_s = new_s.replace(':', '')
            new_s = new_s.replace('：', '')
            if ('票据号码' in new_s and new_s != '票据号码') or (len(new_s) == 30 and new_s.isnumeric()):
                self.data['draftNo'] = new_s.replace('票据号码', '')
            if ('票据状态' in new_s and new_s != '票据状态'):
                self.data['draftStatus'] = new_s.replace('票据状态', '')
            if ('出票日期' in new_s and new_s != '出票日期'):
                self.data['draftDate'] = new_s.replace('出票日期', '')
            if ('汇票到期日' in new_s and new_s != '汇票到期日'):
                self.data['expireDate'] = new_s.replace('汇票到期日', '')
    
    def ocr_predict(self, imgs, path=False):
        '''img is cv2 image object'''
        try:
            ocr = PaddleOCR(cls_thresh=0, use_gpu=False, 
                            det_model_dir='./inference/ch_PP-OCRv2_det_infer', 
                            rec_model_dir='./inference/ch_PP-OCRv2_rec_infer', 
                            cls_model_dir='./inference/ch_ppocr_mobile_v2.0_cls_infer')
            res_ls = []
            index = 0
            if path:
                for img_path in imgs:
                    file_path = os.path.join('result', img_path)
                    img = cv2.imread(file_path)
                    result = ocr.ocr(img)
                    s = ''
                    for line in result:
                        s += line[1][0]
                    res_ls.append(s)
                return res_ls
            if len(imgs) == 1:
                try:
                    result = ocr.ocr(imgs[0])
                    for line in result:
                        res_ls.append(line[1][0])
                except Exception as e:
                    logging.exception(e)
                return res_ls
            for img in imgs:
                # file_path = os.path.join('result', '%s.png' % index)
                # cv2.imwrite(file_path, img)
                try:
                    result = ocr.ocr(img)
                    s = ''
                    for line in result:
                        s += line[1][0]
                except Exception as e:
                    logging.exception(e)
                    s = ''
                res_ls.append(s)
                index += 1
            return res_ls
        except Exception as e:
            logging.exception(e)
            return []


def ocr_predict(imgs, path=False, ocr=None, file_name='', label=False):
    '''img is cv2 image object'''
    # current_dir = os.getcwd()
    # result_file = os.path.join(current_dir, 'ocr_result.txt')
    # ocr_res = {}
    try:
        if not ocr:
            ocr = PaddleOCR(cls_thresh=0, use_gpu=False, 
                            det_model_dir='./inference/test_det', 
                            rec_model_dir='./inference/ch_PP-OCRv2_rec_infer', 
                            cls_model_dir='./inference/ch_ppocr_mobile_v2.0_cls_infer')
        res_ls = []
        index = 0
        if path:
            for img_path in imgs:
                file_path = os.path.join('result', img_path)
                img = cv2.imread(file_path)
                result = ocr.ocr(img)
                s = ''
                for line in result:
                    s += line[1][0]
                res_ls.append(s)
            return res_ls
        if len(imgs) == 1:
            try:
                result = ocr.ocr(imgs[0])
                for line in result:
                    res_ls.append(line[1][0])
            except Exception as e:
                logging.exception(e)
            return res_ls
        for img in imgs:
            # file_path = os.path.join('result', '%s.png' % index)
            # cv2.imwrite(file_path, img)
            if img.shape[0] <= 40 and img.shape[1] <= 40:
                s = ''
                res_ls.append(s)
                index += 1
                continue
            try:
                result = ocr.ocr(img)
                s = ''
                for line in result:
                    s += line[1][0]
            except Exception as e:
                logging.exception(e)
                s = ''
            res_ls.append(s)
            index += 1
        # print('8' * 128)
        # if label:
        #     ocr_res['%s-%s' % (file_name, 'label')] = res_ls
        # else:
        #     ocr_res[file_name] = res_ls
        # with open(result_file, 'w', encoding='utf8') as fp:
        #     fp.write(json.dumps(ocr_res, ensure_ascii=False))
        return res_ls
    except Exception as e:
        logging.exception(e)
        return []


if __name__ == '__main__':
    ls = os.listdir('result')
    ls = sorted(ls, key=lambda s: int(s.split('.')[0]))
    res = ocr_predict(ls, path=True)
    for index, val in enumerate(res):
        print(index, val)